{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/hongtao945/python_code/pathml-master/examples',\n",
       " '/home/hongtao945/anaconda3/envs/py38/lib/python38.zip',\n",
       " '/home/hongtao945/anaconda3/envs/py38/lib/python3.8',\n",
       " '/home/hongtao945/anaconda3/envs/py38/lib/python3.8/lib-dynload',\n",
       " '',\n",
       " '/home/hongtao945/anaconda3/envs/py38/lib/python3.8/site-packages',\n",
       " '..',\n",
       " '../',\n",
       " '../convformer/',\n",
       " '..',\n",
       " '../',\n",
       " '../convformer/']"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('..')\n",
    "sys.path.append('../')\n",
    "sys.path.append('../convformer/')\n",
    "sys.path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import albumentations as A\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import copy\n",
    "import torch\n",
    "from pathml.datasets.MICCAIds import MICCAIDataSet\n",
    "from pathml.ml.hovernet import HoVerNet as HC\n",
    "from torch.optim.lr_scheduler import StepLR\n",
    "from pathml.ml.hovernet_attention import HoVerNet as HA, loss_hovernet, post_process_batch_hovernet\n",
    "from pathml.ml.utils import wrap_transform_multichannel, dice_score\n",
    "from pathml.utils import plot_segmentation\n",
    "from pathml.ml.hovernet_convformer import HoVerNet_ConvFormer as HCF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(30, 10)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hover_transform = A.Compose(\n",
    "    [A.VerticalFlip(p=0.5), \n",
    "     A.HorizontalFlip(p=0.5),\n",
    "     A.RandomRotate90(p=0.5),\n",
    "     A.GaussianBlur(p=0.5),\n",
    "     A.MedianBlur(p=0.5, blur_limit=5)]\n",
    ")\n",
    "train_dataset = MICCAIDataSet('../data/MICCAI2018MoNuSeg', transform=hover_transform)\n",
    "valid_dataset = MICCAIDataSet('../data/MICCAI2018MoNuSeg', transform=None, mode='val')\n",
    "train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "valid_dataloader = torch.utils.data.DataLoader(valid_dataset, batch_size=1, shuffle=False)\n",
    "len(train_dataloader), len(valid_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hongtao945/anaconda3/envs/py38/lib/python3.8/site-packages/torch/functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at /opt/conda/conda-bld/pytorch_1699449229234/work/aten/src/ATen/native/TensorShape.cpp:3526.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "hovernet = HCF(n_channels=3, imgsize=512).to(device)\n",
    "opt = torch.optim.Adam(hovernet.parameters(), lr = 1e-4)\n",
    "scheduler = StepLR(opt, step_size=25, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/75 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|▏         | 1/75 [32:27<40:02:08, 1947.68s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/75:\n",
      "\ttraining loss: 2.5551\tvalidation loss: 2.5893\n",
      "\ttraining dice: 0.4135\tvalidation dice: 0.1888\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  3%|▎         | 2/75 [1:05:04<39:36:31, 1953.31s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/75:\n",
      "\ttraining loss: 1.7853\tvalidation loss: 2.4309\n",
      "\ttraining dice: 0.4895\tvalidation dice: 0.3316\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 3/75 [1:38:17<39:25:33, 1971.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3/75:\n",
      "\ttraining loss: 1.7411\tvalidation loss: 3.2412\n",
      "\ttraining dice: 0.5058\tvalidation dice: 0.1938\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  5%|▌         | 4/75 [2:11:24<38:59:59, 1977.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4/75:\n",
      "\ttraining loss: 1.8032\tvalidation loss: 2.5009\n",
      "\ttraining dice: 0.4955\tvalidation dice: 0.0537\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 5/75 [2:44:13<38:23:35, 1974.50s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5/75:\n",
      "\ttraining loss: 1.6363\tvalidation loss: 2.1478\n",
      "\ttraining dice: 0.5097\tvalidation dice: 0.4728\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 6/75 [3:50:15<50:47:41, 2650.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6/75:\n",
      "\ttraining loss: 1.5727\tvalidation loss: 2.3235\n",
      "\ttraining dice: 0.5725\tvalidation dice: 0.2944\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  9%|▉         | 7/75 [6:02:25<82:39:42, 4376.22s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7/75:\n",
      "\ttraining loss: 1.5734\tvalidation loss: 2.169\n",
      "\ttraining dice: 0.5109\tvalidation dice: 0.3931\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 8/75 [6:41:12<69:18:18, 3723.86s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8/75:\n",
      "\ttraining loss: 1.5386\tvalidation loss: 2.2165\n",
      "\ttraining dice: 0.5664\tvalidation dice: 0.351\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 9/75 [7:20:00<60:16:18, 3287.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9/75:\n",
      "\ttraining loss: 1.5399\tvalidation loss: 2.7603\n",
      "\ttraining dice: 0.5356\tvalidation dice: 0.0416\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 10/75 [7:58:48<54:00:27, 2991.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/75:\n",
      "\ttraining loss: 1.5173\tvalidation loss: 2.0952\n",
      "\ttraining dice: 0.5343\tvalidation dice: 0.459\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▍        | 11/75 [8:37:34<49:33:28, 2787.64s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11/75:\n",
      "\ttraining loss: 1.5256\tvalidation loss: 2.6915\n",
      "\ttraining dice: 0.5328\tvalidation dice: 0.0154\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 12/75 [9:16:21<46:19:58, 2647.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12/75:\n",
      "\ttraining loss: 1.511\tvalidation loss: 2.522\n",
      "\ttraining dice: 0.5767\tvalidation dice: 0.39\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 17%|█▋        | 13/75 [9:55:09<43:55:39, 2550.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13/75:\n",
      "\ttraining loss: 1.5454\tvalidation loss: 2.4075\n",
      "\ttraining dice: 0.531\tvalidation dice: 0.1813\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 19%|█▊        | 14/75 [10:33:57<42:04:47, 2483.40s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14/75:\n",
      "\ttraining loss: 1.4715\tvalidation loss: 2.1515\n",
      "\ttraining dice: 0.598\tvalidation dice: 0.5372\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 15/75 [11:32:34<46:35:03, 2795.05s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15/75:\n",
      "\ttraining loss: 1.5064\tvalidation loss: 2.5447\n",
      "\ttraining dice: 0.4716\tvalidation dice: 0.3939\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 21%|██▏       | 16/75 [12:11:26<43:31:17, 2655.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 16/75:\n",
      "\ttraining loss: 1.4938\tvalidation loss: 2.3282\n",
      "\ttraining dice: 0.5732\tvalidation dice: 0.105\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 23%|██▎       | 17/75 [12:50:13<41:11:36, 2556.84s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 17/75:\n",
      "\ttraining loss: 1.4811\tvalidation loss: 2.2504\n",
      "\ttraining dice: 0.5925\tvalidation dice: 0.1812\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 18/75 [13:29:00<39:23:29, 2487.89s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18/75:\n",
      "\ttraining loss: 1.4893\tvalidation loss: 2.4476\n",
      "\ttraining dice: 0.5523\tvalidation dice: 0.0258\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 25%|██▌       | 19/75 [14:07:47<37:56:58, 2439.61s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19/75:\n",
      "\ttraining loss: 1.4871\tvalidation loss: 2.4827\n",
      "\ttraining dice: 0.5678\tvalidation dice: 0.0744\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 20/75 [14:46:34<36:45:06, 2405.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20/75:\n",
      "\ttraining loss: 1.4712\tvalidation loss: 2.3627\n",
      "\ttraining dice: 0.6108\tvalidation dice: 0.3409\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 21/75 [15:25:19<35:43:26, 2381.59s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 21/75:\n",
      "\ttraining loss: 1.4513\tvalidation loss: 4.9882\n",
      "\ttraining dice: 0.5851\tvalidation dice: 0.0139\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▉       | 22/75 [16:04:06<34:49:03, 2364.97s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 22/75:\n",
      "\ttraining loss: 1.43\tvalidation loss: 4.4149\n",
      "\ttraining dice: 0.5766\tvalidation dice: 0.0\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███       | 23/75 [16:42:56<34:00:42, 2354.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 23/75:\n",
      "\ttraining loss: 1.4071\tvalidation loss: 2.0692\n",
      "\ttraining dice: 0.5768\tvalidation dice: 0.4918\n",
      "val\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 24/75 [17:21:44<33:14:37, 2346.63s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 24/75:\n",
      "\ttraining loss: 1.4311\tvalidation loss: 3.0883\n",
      "\ttraining dice: 0.5734\tvalidation dice: 0.1497\n"
     ]
    }
   ],
   "source": [
    "n_epochs = 75\n",
    "\n",
    "# print performance metrics every n epochs\n",
    "print_every_n_epochs = True\n",
    "\n",
    "# evaluating performance on a random subset of validation mini-batches\n",
    "# this saves time instead of evaluating on the entire validation set\n",
    "n_minibatch_valid = 10\n",
    "\n",
    "epoch_train_losses = {}\n",
    "epoch_valid_losses = {}\n",
    "epoch_train_dice = {}\n",
    "epoch_valid_dice = {}\n",
    "\n",
    "best_epoch = 0\n",
    "\n",
    "for i in tqdm(range(n_epochs)):\n",
    "    minibatch_train_losses = []\n",
    "    minibatch_train_dice = []\n",
    "    \n",
    "    # put model in training mode\n",
    "    hovernet.train()\n",
    "    \n",
    "    for data in train_dataloader:\n",
    "        # send the data to the GPU\n",
    "        images = data[0].float().to(device)\n",
    "        masks = data[1].to(device)\n",
    "        hv = data[2].float().to(device)\n",
    "        \n",
    "        # zero out gradient\n",
    "        opt.zero_grad()\n",
    "        \n",
    "        # forward pass\n",
    "        outputs = hovernet(images)\n",
    "        \n",
    "        # compute loss\n",
    "        loss = loss_hovernet(outputs = outputs, ground_truth = [masks, hv])\n",
    "        \n",
    "        # track loss\n",
    "        minibatch_train_losses.append(loss.item())\n",
    "        \n",
    "        # also track dice score to measure performance\n",
    "        preds_detection = post_process_batch_hovernet(outputs, n_classes=None)\n",
    "        truth_binary = masks[:, -1, :, :] == 0\n",
    "        dice = dice_score(preds_detection, truth_binary.cpu().numpy())\n",
    "        minibatch_train_dice.append(dice)\n",
    "        \n",
    "        # compute gradients\n",
    "        loss.backward()\n",
    "        \n",
    "        # step optimizer and scheduler\n",
    "        opt.step()\n",
    "    \n",
    "    #step LR scheduler\n",
    "    scheduler.step()\n",
    "    print('val')\n",
    "    # evaluate on random subset of validation data\n",
    "    hovernet.eval()\n",
    "    minibatch_valid_losses = []\n",
    "    minibatch_valid_dice = []\n",
    "    # randomly choose minibatches for evaluating\n",
    "    minibatch_ix = np.random.choice(range(len(valid_dataloader)), replace=False, size=n_minibatch_valid)\n",
    "    with torch.no_grad():\n",
    "        for j, data in enumerate(valid_dataloader):\n",
    "            if j in minibatch_ix:\n",
    "                # send the data to the GPU\n",
    "                images = data[0].float().to(device)\n",
    "                masks = data[1].to(device)\n",
    "                hv = data[2].float().to(device)\n",
    "\n",
    "                # forward pass\n",
    "                outputs = hovernet(images)\n",
    "\n",
    "                # compute loss\n",
    "                loss = loss_hovernet(outputs = outputs, ground_truth = [masks, hv], n_classes=None)\n",
    "\n",
    "                # track loss\n",
    "                minibatch_valid_losses.append(loss.item())\n",
    "\n",
    "                # also track dice score to measure performance\n",
    "                preds_detection = post_process_batch_hovernet(outputs, n_classes=None)\n",
    "                truth_binary = masks[:, -1, :, :] == 0\n",
    "                dice = dice_score(preds_detection, truth_binary.cpu().numpy())\n",
    "                minibatch_valid_dice.append(dice)\n",
    "    \n",
    "    # average performance metrics over minibatches\n",
    "    mean_train_loss = np.mean(minibatch_train_losses)\n",
    "    mean_valid_loss = np.mean(minibatch_valid_losses)\n",
    "    mean_train_dice = np.mean(minibatch_train_dice)\n",
    "    mean_valid_dice = np.mean(minibatch_valid_dice)\n",
    "    \n",
    "    # save the model with best performance\n",
    "    if i != 0:\n",
    "        if mean_valid_loss < min(epoch_valid_losses.values()):\n",
    "            best_epoch = i\n",
    "            torch.save(hovernet.state_dict(), f\"hovernet_convformer_withskip_best_perf.pt\")\n",
    "    \n",
    "    # track performance over training epochs\n",
    "    epoch_train_losses.update({i : mean_train_loss})\n",
    "    epoch_valid_losses.update({i : mean_valid_loss})\n",
    "    epoch_train_dice.update({i : mean_train_dice})\n",
    "    epoch_valid_dice.update({i : mean_valid_dice})\n",
    "    \n",
    "    if print_every_n_epochs is not None:\n",
    "        if i % print_every_n_epochs == print_every_n_epochs - 1:\n",
    "            print(f\"Epoch {i+1}/{n_epochs}:\")\n",
    "            print(f\"\\ttraining loss: {np.round(mean_train_loss, 4)}\\tvalidation loss: {np.round(mean_valid_loss, 4)}\")\n",
    "            print(f\"\\ttraining dice: {np.round(mean_train_dice, 4)}\\tvalidation dice: {np.round(mean_valid_dice, 4)}\")\n",
    "\n",
    "# save fully trained model\n",
    "torch.save(hovernet.state_dict(), f\"hovernet_convformer_withskip_fully_trained.pt\")\n",
    "print(f\"\\nEpoch with best validation performance: {best_epoch}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
